import torch
# 图表和可视化
import matplotlib.pyplot as plt

def visualize_predictions(model, device, test_loader, num_samples=5):
    model.eval()
    fig, axes = plt.subplots(1, num_samples, figsize=(15, 3))
    
    with torch.no_grad():
        data, target = next(iter(test_loader))
        data, target = data.to(device), target.to(device)
        output = model(data)
        pred = output.argmax(dim=1, keepdim=True)
        
        for i in range(num_samples):
            img = data[i].cpu().squeeze().numpy()
            axes[i].imshow(img, cmap='gray')
            axes[i].set_title(f'Pred: {pred[i].item()}, True: {target[i].item()}')
            axes[i].axis('off')
    
    plt.tight_layout()
    plt.show()
